#ifndef _COMM_H
#define _COMM_H

#include <openssl/evp.h>
#include <openssl/rsa.h>
#include <openssl/pem.h>
#include <string.h>

#ifdef DEBUG
#define print(format, ...)                                                     \
    {                                                                          \
        fprintf(stderr, "[file:%s func:%s line:%d]. ", __FILE__, __FUNCTION__, \
                __LINE__);                                                     \
        fprintf(stderr, format, ##__VA_ARGS__);                                \
    }
#else
#define print(format, ...)
#endif

static const unsigned char kExampleRSAKeyDER[] = {
    0x30, 0x82, 0x02, 0x5c, 0x02, 0x01, 0x00, 0x02, 0x81, 0x81, 0x00, 0xf8,
    0xb8, 0x6c, 0x83, 0xb4, 0xbc, 0xd9, 0xa8, 0x57, 0xc0, 0xa5, 0xb4, 0x59,
    0x76, 0x8c, 0x54, 0x1d, 0x79, 0xeb, 0x22, 0x52, 0x04, 0x7e, 0xd3, 0x37,
    0xeb, 0x41, 0xfd, 0x83, 0xf9, 0xf0, 0xa6, 0x85, 0x15, 0x34, 0x75, 0x71,
    0x5a, 0x84, 0xa8, 0x3c, 0xd2, 0xef, 0x5a, 0x4e, 0xd3, 0xde, 0x97, 0x8a,
    0xdd, 0xff, 0xbb, 0xcf, 0x0a, 0xaa, 0x86, 0x92, 0xbe, 0xb8, 0x50, 0xe4,
    0xcd, 0x6f, 0x80, 0x33, 0x30, 0x76, 0x13, 0x8f, 0xca, 0x7b, 0xdc, 0xec,
    0x5a, 0xca, 0x63, 0xc7, 0x03, 0x25, 0xef, 0xa8, 0x8a, 0x83, 0x58, 0x76,
    0x20, 0xfa, 0x16, 0x77, 0xd7, 0x79, 0x92, 0x63, 0x01, 0x48, 0x1a, 0xd8,
    0x7b, 0x67, 0xf1, 0x52, 0x55, 0x49, 0x4e, 0xd6, 0x6e, 0x4a, 0x5c, 0xd7,
    0x7a, 0x37, 0x36, 0x0c, 0xde, 0xdd, 0x8f, 0x44, 0xe8, 0xc2, 0xa7, 0x2c,
    0x2b, 0xb5, 0xaf, 0x64, 0x4b, 0x61, 0x07, 0x02, 0x03, 0x01, 0x00, 0x01,
    0x02, 0x81, 0x80, 0x74, 0x88, 0x64, 0x3f, 0x69, 0x45, 0x3a, 0x6d, 0xc7,
    0x7f, 0xb9, 0xa3, 0xc0, 0x6e, 0xec, 0xdc, 0xd4, 0x5a, 0xb5, 0x32, 0x85,
    0x5f, 0x19, 0xd4, 0xf8, 0xd4, 0x3f, 0x3c, 0xfa, 0xc2, 0xf6, 0x5f, 0xee,
    0xe6, 0xba, 0x87, 0x74, 0x2e, 0xc7, 0x0c, 0xd4, 0x42, 0xb8, 0x66, 0x85,
    0x9c, 0x7b, 0x24, 0x61, 0xaa, 0x16, 0x11, 0xf6, 0xb5, 0xb6, 0xa4, 0x0a,
    0xc9, 0x55, 0x2e, 0x81, 0xa5, 0x47, 0x61, 0xcb, 0x25, 0x8f, 0xc2, 0x15,
    0x7b, 0x0e, 0x7c, 0x36, 0x9f, 0x3a, 0xda, 0x58, 0x86, 0x1c, 0x5b, 0x83,
    0x79, 0xe6, 0x2b, 0xcc, 0xe6, 0xfa, 0x2c, 0x61, 0xf2, 0x78, 0x80, 0x1b,
    0xe2, 0xf3, 0x9d, 0x39, 0x2b, 0x65, 0x57, 0x91, 0x3d, 0x71, 0x99, 0x73,
    0xa5, 0xc2, 0x79, 0x20, 0x8c, 0x07, 0x4f, 0xe5, 0xb4, 0x60, 0x1f, 0x99,
    0xa2, 0xb1, 0x4f, 0x0c, 0xef, 0xbc, 0x59, 0x53, 0x00, 0x7d, 0xb1, 0x02,
    0x41, 0x00, 0xfc, 0x7e, 0x23, 0x65, 0x70, 0xf8, 0xce, 0xd3, 0x40, 0x41,
    0x80, 0x6a, 0x1d, 0x01, 0xd6, 0x01, 0xff, 0xb6, 0x1b, 0x3d, 0x3d, 0x59,
    0x09, 0x33, 0x79, 0xc0, 0x4f, 0xde, 0x96, 0x27, 0x4b, 0x18, 0xc6, 0xd9,
    0x78, 0xf1, 0xf4, 0x35, 0x46, 0xe9, 0x7c, 0x42, 0x7a, 0x5d, 0x9f, 0xef,
    0x54, 0xb8, 0xf7, 0x9f, 0xc4, 0x33, 0x6c, 0xf3, 0x8c, 0x32, 0x46, 0x87,
    0x67, 0x30, 0x7b, 0xa7, 0xac, 0xe3, 0x02, 0x41, 0x00, 0xfc, 0x2c, 0xdf,
    0x0c, 0x0d, 0x88, 0xf5, 0xb1, 0x92, 0xa8, 0x93, 0x47, 0x63, 0x55, 0xf5,
    0xca, 0x58, 0x43, 0xba, 0x1c, 0xe5, 0x9e, 0xb6, 0x95, 0x05, 0xcd, 0xb5,
    0x82, 0xdf, 0xeb, 0x04, 0x53, 0x9d, 0xbd, 0xc2, 0x38, 0x16, 0xb3, 0x62,
    0xdd, 0xa1, 0x46, 0xdb, 0x6d, 0x97, 0x93, 0x9f, 0x8a, 0xc3, 0x9b, 0x64,
    0x7e, 0x42, 0xe3, 0x32, 0x57, 0x19, 0x1b, 0xd5, 0x6e, 0x85, 0xfa, 0xb8,
    0x8d, 0x02, 0x41, 0x00, 0xbc, 0x3d, 0xde, 0x6d, 0xd6, 0x97, 0xe8, 0xba,
    0x9e, 0x81, 0x37, 0x17, 0xe5, 0xa0, 0x64, 0xc9, 0x00, 0xb7, 0xe7, 0xfe,
    0xf4, 0x29, 0xd9, 0x2e, 0x43, 0x6b, 0x19, 0x20, 0xbd, 0x99, 0x75, 0xe7,
    0x76, 0xf8, 0xd3, 0xae, 0xaf, 0x7e, 0xb8, 0xeb, 0x81, 0xf4, 0x9d, 0xfe,
    0x07, 0x2b, 0x0b, 0x63, 0x0b, 0x5a, 0x55, 0x90, 0x71, 0x7d, 0xf1, 0xdb,
    0xd9, 0xb1, 0x41, 0x41, 0x68, 0x2f, 0x4e, 0x39, 0x02, 0x40, 0x5a, 0x34,
    0x66, 0xd8, 0xf5, 0xe2, 0x7f, 0x18, 0xb5, 0x00, 0x6e, 0x26, 0x84, 0x27,
    0x14, 0x93, 0xfb, 0xfc, 0xc6, 0x0f, 0x5e, 0x27, 0xe6, 0xe1, 0xe9, 0xc0,
    0x8a, 0xe4, 0x34, 0xda, 0xe9, 0xa2, 0x4b, 0x73, 0xbc, 0x8c, 0xb9, 0xba,
    0x13, 0x6c, 0x7a, 0x2b, 0x51, 0x84, 0xa3, 0x4a, 0xe0, 0x30, 0x10, 0x06,
    0x7e, 0xed, 0x17, 0x5a, 0x14, 0x00, 0xc9, 0xef, 0x85, 0xea, 0x52, 0x2c,
    0xbc, 0x65, 0x02, 0x40, 0x51, 0xe3, 0xf2, 0x83, 0x19, 0x9b, 0xc4, 0x1e,
    0x2f, 0x50, 0x3d, 0xdf, 0x5a, 0xa2, 0x18, 0xca, 0x5f, 0x2e, 0x49, 0xaf,
    0x6f, 0xcc, 0xfa, 0x65, 0x77, 0x94, 0xb5, 0xa1, 0x0a, 0xa9, 0xd1, 0x8a,
    0x39, 0x37, 0xf4, 0x0b, 0xa0, 0xd7, 0x82, 0x27, 0x5e, 0xae, 0x17, 0x17,
    0xa1, 0x1e, 0x54, 0x34, 0xbf, 0x6e, 0xc4, 0x8e, 0x99, 0x5d, 0x08, 0xf1,
    0x2d, 0x86, 0x9d, 0xa5, 0x20, 0x1b, 0xe5, 0xdf,
};

// 生成密钥对
static EVP_PKEY* load_example_rsa_key(void)
{
    EVP_PKEY* ret             = NULL;
    const unsigned char* derp = kExampleRSAKeyDER;
    EVP_PKEY* pkey            = NULL;
    RSA* rsa                  = NULL;
    d2i_RSAPrivateKey(&rsa, &derp, sizeof(kExampleRSAKeyDER));

    pkey = EVP_PKEY_new();
    EVP_PKEY_set1_RSA(pkey, rsa);

    ret  = pkey;
    pkey = NULL;
    EVP_PKEY_free(pkey);
    RSA_free(rsa);
    return ret;
}

// 生成公钥和私钥文件
static int generate_key_file(const char* pubkey_file, const char* prikey_file,
                      unsigned char* passwd)
{
    RSA* rsa = RSA_new();
    if (rsa == NULL) {
        print("RSA_new failed .\n");
        goto err;
    }
    BIGNUM* bn = BN_new();
    if (bn == NULL) {
        print("BN_new failed.\n");
        goto err;
    }

    if (!BN_set_word(bn, RSA_F4)) {
        print("BN_set_word failed.\n");
        goto err;
    }
    if (!RSA_generate_key_ex(rsa, 1024, bn, NULL)) {
        print("RSA_generate_key_ex failed .\n");
        goto err;
    }

    FILE* fpubkey = fopen(pubkey_file, "w");
    if (fpubkey == NULL) {
        print("open %s failed .\n", pubkey_file);
        goto err;
    }
    PEM_write_RSAPublicKey(fpubkey, rsa);

    FILE* fprikey = fopen(prikey_file, "w");
    if (fprikey == NULL) {
        print("open %s failed .\n", prikey_file);
        goto err;
    }
    PEM_write_RSAPrivateKey(fprikey, rsa, EVP_des_ede3_ofb(), passwd,
                            strlen((const char*)passwd), NULL, NULL);
    RSA_free(rsa);
    BN_free(bn);
    fclose(fpubkey);
    fclose(fprikey);
    return 0;
err:
    if (rsa) RSA_free(rsa);
    if (bn) BN_free(bn);
    if (fpubkey) fclose(fpubkey);
    if (fprikey) fclose(fprikey);
    return 1;
}

// 生成公钥文件和私钥文件，私钥文件带密码
static const int RSA_KEY_LENGTH = 1024;
int generate_key_files(const char *pub_keyfile, const char *pri_keyfile, 
		unsigned char *passwd)
{
	RSA* rsa = RSA_new();
	BIGNUM* bn = BN_new();
	BN_set_word(bn, RSA_F4);
	
	RSA_generate_key_ex(rsa, RSA_KEY_LENGTH, bn, NULL);

	FILE* fpub = fopen(pub_keyfile, "w");
	PEM_write_RSAPublicKey(fpub, rsa);
	fclose(fpub);

	FILE* fpri = fopen(pri_keyfile, "w");
	PEM_write_RSAPrivateKey(fpri, rsa, EVP_des_ede3_ofb(),
			passwd, strlen((const char*)passwd), NULL, NULL);
	fclose(fpri);

	RSA_free(rsa);
	BN_free(bn);
	return 0;
}


// 读取公钥
static EVP_PKEY* read_pubkey(const char* pubkey_file)
{
    FILE* fpub = fopen(pubkey_file, "r");
    if (fpub == NULL) {
        print("fopen %s failed.\n", pubkey_file);
        goto err;
    }
    RSA* rsa = PEM_read_RSAPublicKey(fpub, NULL, NULL, NULL);
    // RSA* rsa = PEM_read_RSA_PUBKEY(fpub, NULL, NULL, NULL);
    if (rsa == NULL) {
        print("PEM_read_RSAPublicKey failed . \n");
        goto err;
    }
    EVP_PKEY* key = EVP_PKEY_new();
    if (key == NULL) {
        print("EVP_PKEY_new failed.\n");
        goto err;
    }
    EVP_PKEY_assign(key, EVP_PKEY_RSA, rsa);
    // EVP_PKEY_set1_RSA(key, rsa);
    fclose(fpub);
    return key;
err:
    if (fpub) fclose(fpub);
    exit(1);
}

// 读取私钥
static EVP_PKEY* read_prikey(const char* prikey_file)
{
    FILE* fpri = fopen(prikey_file, "r");
    if (fpri == NULL) {
        print("fopen %s failed.\n", prikey_file);
        goto err;
    }
    RSA* rsa = PEM_read_RSAPrivateKey(fpri, NULL, NULL, NULL);
    if (rsa == NULL) {
        print("PEM_read_RSAPrivateKey failed . \n");
        goto err;
    }
    EVP_PKEY* key = EVP_PKEY_new();
    if (key == NULL) {
        print("EVP_PKEY_new failed.\n");
        goto err;
    }
    EVP_PKEY_assign(key, EVP_PKEY_RSA, rsa);

    fclose(fpri);
    return key;
err:
    if (fpri) fclose(fpri);
    exit(1);
}

// 打开公钥文件，返回EVP_PKEY结构的指针
EVP_PKEY* open_public_key(const char *keyfile)
{
	FILE* fpub = fopen(keyfile, "r");

	RSA* rsa = PEM_read_RSAPublicKey(fpub, NULL, NULL, NULL);
	fclose(fpub);

	EVP_PKEY* pkey = EVP_PKEY_new();
	EVP_PKEY_assign_RSA(pkey, rsa);
	return pkey;
}

// 打开私钥文件，返回EVP_PKEY结构的指针
EVP_PKEY* open_private_key(const char *keyfile)
{
	OpenSSL_add_all_algorithms();

	FILE* fpri = fopen(keyfile, "r");
	RSA* rsa = PEM_read_RSAPrivateKey(fpri, NULL, NULL, NULL);
	fclose(fpri);

	EVP_PKEY* key = EVP_PKEY_new();
	EVP_PKEY_assign_RSA(key, rsa);
	return key;
}


#endif /* _COMM_H */
